#  -*- coding: utf-8 -*-

from pymongo import UpdateOne,ASCENDING, DESCENDING
from factor.base_factor import BaseFactor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes,get_all_indexes_date,calc_negative_diff_dates,get_trading_dates,get_all_codes_trading_date
from util.database import DB_CONN
import time
from datetime import datetime, timedelta
from pandas import DataFrame
import pandas as pd

"""
实现大盘指数的相关因子计算
"""


class MainIndexFactor(BaseFactor):
    def __init__(self):
        BaseFactor.__init__(self, name='main_index')


    def comupute_change_rate_count(self,date):
        #每天涨跌停数据,大于9.5，小于-9.5
        start_time = time.time()
        change_rate_dict = dict()
        change_collection = DB_CONN["change_rate"]

        codes = get_all_codes_trading_date(date)
        codes_num = len(codes)
        dm = DataModule()
        df_daily = dm.get_k_data("999999", index=True,  begin_date=date, end_date=date)
        if df_daily.index.size > 0:
            df_daily.set_index(['date'], 1, inplace=True)
            amount = round(df_daily.loc[date]['amount']/1E8,2)

        factor_cursor = change_collection.find({'date': date,"index":False},projection={'_id': False})
        change_rate_df = DataFrame([x for x in factor_cursor])

        above_9_num = len(change_rate_df[change_rate_df["close_change_rate"] > 9.5])
        below_N9_num = len(change_rate_df[change_rate_df["close_change_rate"] < -9.5])

        change_rate_dict['above_9_num'] = above_9_num
        change_rate_dict['above_9_ratio'] = round(100*above_9_num/codes_num,2)
        change_rate_dict['below_N9_num'] = below_N9_num
        change_rate_dict['below_N9_ratio'] = round(100*below_N9_num/codes_num,2)
        change_rate_dict['amount'] = amount
        end_time = time.time()
        print(f"计算{date}大盘数据,耗时{end_time - start_time} 秒！")
        return change_rate_dict


    def compute(self, begin_date, end_date):
        """
        计算时间段内每日大盘类因子数据，并保存到数据库中
        """
        start_time = time.time()

        self.collection.create_index([('date', 1)])
        dates = get_trading_dates(begin_date,end_date)

        update_requests = []

        for date in dates:
            factor_dict = dict()
            change_rate_dict = self.comupute_change_rate_count(date)
            factor_dict.update(change_rate_dict)
            #print(factor_dict)
            update_requests.append(
                UpdateOne(
                    {'date': date},
                    {'$set': factor_dict},
                    upsert=True))


        if len(update_requests) > 0:
            update_result = self.collection.bulk_write(update_requests, ordered=False)

            end_time = time.time()
            print(f'填充大盘因子数据，插入{update_result.upserted_count}条，更新：{update_result.modified_count}条,'
                  f'耗时：{end_time - start_time}秒',
                  flush=True)

        return


if __name__ == '__main__':
    # 执行因子的提取任务
    #hfq =HfqMAFactor()

    MainIndexFactor().compute('2014-01-01', '2019-05-31')
